
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

class ChatClient:
    def __init__(self, model_path):
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            trust_remote_code=True,
            pad_token='<|endoftext|>',
            local_files_only=True
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path,
            trust_remote_code=True,
            torch_dtype=torch.float16,
            device_map="auto",
            local_files_only=True
        )
        self.chat_history = []

    def generate_response(self, user_input):
        self.chat_history.append({"role": "user", "content": user_input})
        
        # 修正输入处理逻辑
        inputs = self.tokenizer.apply_chat_template(
            self.chat_history,
            add_generation_prompt=True,
            return_tensors="pt"
        ).to(self.model.device)
        
        outputs = self.model.generate(
            input_ids=inputs,
            max_new_tokens=1024,
            pad_token_id=self.tokenizer.pad_token_id,
            do_sample=True,
            temperature=0.7,
            top_p=0.9
        )
        
        response = self.tokenizer.decode(
            outputs[0][inputs.shape[1]:],
            skip_special_tokens=True
        )
        self.chat_history.append({"role": "assistant", "content": response})
        return response

if __name__ == "__main__":
    MODEL_PATH = "/home/fangning/work/LLM/models/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
    client = ChatClient(MODEL_PATH)
    print("DeepSeek对话客户端已启动(输入quit退出)")
    
    try:
        while True:
            user_input = input("\n用户: ")
            if user_input.lower() in ['quit', 'exit']:
                break
            print(f"\nAI: {client.generate_response(user_input)}")
    except KeyboardInterrupt:
        print("\n对话已终止")
